import PIL
from tensorflow.keras.utils import Sequence
import os
import math
from PIL import Image
import numpy as np
import random


class MySeg2dDatasets(Sequence):
    def __init__(self,work_dir,batch_size,input_shape,mode,num_classes):
        self.work_dir = work_dir
        self.batch_size = batch_size
        self.target_size = input_shape[:-1]
        self.mode = mode
        self.num_classes = num_classes

        self.images_dir = os.path.join(work_dir,f'{mode}/JPEGImages')
        self.labels_dir = os.path.join(work_dir,f'{mode}/Segmentations')

        self.file_names = os.listdir(self.labels_dir)
        # for file_name in self.file_names:
        #     print(file_name)
    
    def __len__(self):
        return math.floor(len(self.file_names) / self.batch_size)
    

    def normalize(self,img):
        max_val = np.max(img)
        min_val = np.min(img)
        val_range = max_val - min_val
        norm_0_1 = (img-min_val)/val_range
        img = np.clip(2*norm_0_1-1,-1,1)
        return img


    def __getitem__(self, index):
        batch_names = self.file_names[index * self.batch_size:(index + 1) *
        self.batch_size]

        xs = []
        ys = []
        for batch_name in batch_names:
            x_path = os.path.join(self.images_dir,f'{batch_name[:-4]}.jpg')
            y_path = os.path.join(self.labels_dir,batch_name)

            # --------------------------
            # 图片
            # --------------------------
            x = Image.open(x_path).convert('RGB')
            # resize
            x = x.resize(self.target_size,PIL.Image.Resampling.BILINEAR)
            # 转成array
            x_array = np.array(x)
            # 归一化
            x_array = self.normalize(x_array)
            xs.append(x_array)

            # --------------------------
            # 标签
            # --------------------------
            y = Image.open(y_path)
            # resize
            y = y.resize(self.target_size,PIL.Image.Resampling.NEAREST)
            # 转成array
            y_array = np.array(y)
            # one-hot
            y_onehot = np.zeros(shape=self.target_size+[self.num_classes])
            for i in range(self.num_classes):
                y_onehot[:,:,i] = (y_array==i)
            
            # input(np.unique(y_onehot))
            ys.append(y_onehot)
        
        xs = np.array(xs).astype('float32')
        ys = np.array(ys).astype('float32')
        # ys = np.expand_dims(ys,axis=-1)
        # print(xs.shape,ys.shape)
        # print(np.unique(xs),np.unique(ys))

        return xs,ys

    # def on_epoch_end(self):        
        # randnum = random.randint(0,100)
        # random.seed(randnum)
        # random.shuffle(train_x)
        # random.seed(randnum)
        # random.shuffle(train_y)

        

# train_datasets = MySeg2dDatasets(work_dir="C:/Users/hblee/Documents/datasets/seg2d",
#                                 batch_size=2,
#                                 input_shape=[512,512,3],
#                                 mode='train')

# train_datasets[2]

